#' @title ## 单基因KM展示
#' @TODO 给定的基因 KM曲线展示图
#' @param geneList 需要绘制的基因，字符串向量
#' @param clinical 需要包含sample time status 等表头
#' @param exp 基因在行名，样本在列名的表达谱
#' @param topN 展示基因数目，一般为6或者8
#' @param saveplot 是否在本地保存图片
#' @param surtime_unit 生存时间 时间单位，月份12，天数365
#' @param output_dir 输出结果路径
#' @param var_name 用来命名文件夹与文件
#' @return  list，合成后的图片 与 cutoff结果
#'
#' @author *WYK*
KM_list <- function(exp = NULL, clinical = NULL, geneList = NULL,legend_pos = c(0.85, 0.85),
                    topN = 8, saveplot = F, surtime_unit = 12, output_dir = "./", var_name = NULL,
                    best_cut = F,w = 4.2,h = 4.7) {
  library(cowplot)
  library(tidyverse)
  library(survival)
  library(survminer)

  if (topN > length(geneList)) {
    cli::cli_alert_danger("The num of gene displayed is less than num of geneList.")
  }

  if (length(var_name) == 0) {
    var_name <- paste0(paste0(sample(c(letters, LETTERS), 4), collapse = ""))
  }

  clinical <- clinical %>%
    select(sample, time, status)

  exp_use <- exp[geneList, ] %>%
    na.omit() %>% 
    t() %>%
    as.data.frame()

  kmdata <- inner_join(clinical, exp_use %>% rownames_to_column(var = "sample"))

  data_tmp <- kmdata %>% dplyr::select(-sample, -time, -status)

  i <- 1
  a <- list()
  n <- ncol(data_tmp)
  cutoff_table <- data.frame()
  repeat{
    if (best_cut) {
      sur.cut <- surv_cutpoint(
        data = kmdata, time = "time", event = "status",
        variables = kmdata %>%
          select(-c(sample, time, status)) %>%
          colnames(.) %>% .[i]
      )

      cut <- summary(sur.cut)$cutpoint

      a[[i]] <- factor(ifelse(data_tmp[, i] >= cut, "High", "Low"), levels = c("Low", "High"))
    } else {
      a[[i]] <- factor(ifelse(data_tmp[, i] <= median(data_tmp[, i]), "Low", "High"), levels = c("Low", "High"))
    }

    cutoff_table[i, 1] <- colnames(data_tmp)[i]
    cutoff_table[i, 2] <- round(median(data_tmp[, i]), 3)

    names(a)[i] <- paste0(gsub("-", "_", colnames(data_tmp)[i]), "_group")
    data_tmp[, names(a)[i]] <- a[[i]]
    i <- i + 1
    if (i > n) {
      break
    }
  } # 构建分组文件，以表达量中位数为高低分组

  colnames(cutoff_table) <- c("gene", "cutoff")

  vars_for_table <- colnames(data_tmp)[(n + 1):(2 * n)]

  kmdata <- inner_join(kmdata, data_tmp)

  formulas <- sapply(vars_for_table, function(x) {
    as.formula(paste0("Surv(time, status) ~ ", x))
  })

  fits <- lapply(formulas, function(x) {
    surv_fit(x, data = kmdata)
  })

  univ_models <- lapply(formulas, function(x) {
    coxph(x, data = kmdata)
  })

  coxtmp <- lapply(univ_models, function(x) {
    summary(x)
  })

  tmp <- lapply(coxtmp, function(x) {
    logrank_pvalue <- x$sctest["pvalue"]
  })

  tmp_name <- unlist(tmp) %>%
    sort() %>%
    .[1:topN] %>%
    names(.) %>%
    str_split(., pattern = "\\.pvalue", simplify = T) %>%
    .[, 1]

  coxtmp <- coxtmp[tmp_name] # 筛选topN个单因素cox结果文件

  p_chara <- lapply(coxtmp, function(x) {

    HR <- x$coefficients[2] %>% as.numeric()
    logrank_pvalue <- x$sctest["pvalue"]
    lower_.95 <- x$conf.int[, "lower .95"] 
    upper_.95 <- x$conf.int[, "upper .95"] 
    C <- x$concordance[1]

    p_chara <- paste0(
      ifelse(logrank_pvalue < 0.001, "P < 0.001", paste0("P = ", round(logrank_pvalue, 3))),
      "\n",
      "HR = ", round(HR, 2),
      "\n95% CI = ", round(lower_.95, 2), " - ", round(upper_.95, 2)
      # "\nC-index = ", round(C, 2)
    )
  }) # 生成单因素cox HR结果信息

  legend_title <- str_split(tmp_name, pattern = "_", simplify = T) %>%
    .[, 1] %>%
    gsub("_", "-", .) %>%
    as.list()

  # legend_labs <- map(unlist(legend_title), function(x) {
  #   value <- cutoff_table %>%
  #     filter(gene == x) %>%
  #     pull(cutoff)
  #   c(paste0("[<=]",value),paste0("[>]",value)) %>% as.character()
  # })

  legend_labs <- map(unlist(legend_title), function(x) {
    c('Low','High')
  })

  names(legend_labs) <- unlist(legend_title)
  names(fits) <- str_remove_all(names(fits),"_group")
  names(p_chara) <- str_remove_all(names(p_chara),"_group")
  names(legend_title) <- unlist(legend_title)
  
  if (surtime_unit == 12) {
    xlab_chara <- "Time in months"
  } else if (surtime_unit == 365) {
    xlab_chara <- "Time in days"
  } else {
    xlab_chara <- "Time"
  }

  Figure_kmlist <- lapply(unlist(legend_title), function(i) {
    ggsurvplot(
      fits[[i]],
      data = kmdata,
      # color = "strata",
      surv.median.line = "hv",
      legend.title = legend_title[[i]],
      legend.labs = legend_labs[[i]],
      # risk.table.title = "Expression group",
      palette = c("#377EB8", "#E41A1C"),
      ggtheme = survminer::theme_survminer(14),
      pval = p_chara[[i]], # 
      pval.size = 4.2,
      risk.table = T,
      legend = legend_pos,
      # tables.theme = egg::theme_article(14),
      xlab = xlab_chara,
      # legend = "top",
      conf.int = TRUE
    )
  })

  i <- 1
  kmlist <- list()
  repeat{
    kmlist[[i]] <- cowplot::plot_grid(
      Figure_kmlist[[i]]$plot + theme(axis.title.x.bottom = element_blank()),
      Figure_kmlist[[i]]$table ,
      ncol = 1,
      align = "v",
      rel_heights = c(0.7, 0.3)
    )
    i <- i + 1
    if (i > topN) {
      break
    }
  }

  names(kmlist) <- legend_title %>% unlist()

  nrow_used <- case_when(
    between(topN, 1, 4) ~ 1,
    between(topN, 5, 8) ~ 2,
    between(topN, 9, 12) ~ 3,
    between(topN, 13, 16) ~ 4,
    between(topN, 17, 20) ~ 5,
    topN > 21 ~ 6
  )

  width_used <- case_when(
    topN < 5 ~ case_when(
      topN == 1 ~ w * 1,
      topN == 2 ~ w * 2,
      topN == 3 ~ w * 3,
      topN == 4 ~ w * 4
    ),
    topN > 4 ~ w * 4
  )

  # 根据图片排列行数 判断出图高度
  height_used <- switch(nrow_used,
    h,
    h * 2,
    h * 3,
    h * 4,
    h * 5,
    h * 6
  )

  Figure_kmlist_1 <- plot_grid(
    plotlist = kmlist,
    nrow = nrow_used
    # labels = "AUTO",
    # label_size = 22
  )

  if (saveplot) {
    if (!dir.exists(sprintf("%s/single_gene_KMcurve/", output_dir))) {
      dir.create(sprintf("%s/single_gene_KMcurve/", output_dir), recursive = T)
    }

    write.table(
      x = cutoff_table, file = sprintf("%s/single_gene_KMcurve/KM_CutoffTable_%s.txt", output_dir, var_name),
      quote = F, row.names = F, sep = "\t"
    )

    ggsave2(
      filename = sprintf("%s/single_gene_KMcurve/Figure_KM_AllGene_%s.pdf", output_dir, var_name),
      plot = Figure_kmlist_1, width = width_used, height = height_used
    )
    # ggsave2(
    #   filename = sprintf("%soutput/single_gene_KMcurve/Figure_KM_LIST_%s.tiff", output_dir, var_name),
    #   plot = Figure_kmlist_1, width = width_used, height = height_used, dpi = 300
    # )
    # ggsave2(
    #   filename = sprintf("%soutput/single_gene_KMcurve/Figure_KM_LIST_%s_dpi72.tiff", output_dir, var_name),
    #   plot = Figure_kmlist_1, width = width_used, height = height_used, dpi = 72
    # )

    walk(seq_along(kmlist), function(x) {
      ggsave2(
        filename = sprintf("%s/single_gene_KMcurve/Figure_KM_%s_%s.pdf", output_dir, names(kmlist)[x], var_name),
        plot = kmlist[[x]], width = w, height = h
      )
      # ggsave2(
      #   filename = sprintf("%soutput/single_gene_KMcurve/Figure_KM_%s_%s.tiff", output_dir, names(kmlist)[x], var_name),
      #   plot = kmlist[[x]], width = w, height = 5, dpi = 300
      # )
      # ggsave2(
      #   filename = sprintf("%soutput/single_gene_KMcurve/Figure_KM_%s_%s_dpi72.tiff", output_dir, names(kmlist)[x], var_name),
      #   plot = kmlist[[x]], width = w, height = 5, dpi = 72
      # )
    })
  }
  # print(Figure_kmlist_1)
  # print(cutoff_table %>% knitr::kable(digits = 3,align = 'c'))

  X <- list("Figure_KM" = Figure_kmlist_1, "SingleGene_KM" = kmlist, "cutoff_table" = cutoff_table)

  return(X)
}